import torch
import torch.nn as nn
import torch.nn.functional as F


class Model(nn.Module):
    def __init__(self, vocab_size, embedding_size, hidden_size, dropout):
        super(Model, self).__init__()
        self.embedding = nn.Embedding(vocab_size, embedding_size)
        self.bilstm = nn.LSTM(embedding_size, hidden_size, bidirectional=True)
        self.dropout = nn.Dropout(dropout)
        self.fc = nn.Linear(2 * hidden_size, 2)

    def forward(self, x):
        emb = self.embedding(x)
        emb = self.dropout(emb)
        emb = emb.permute(1, 0, 2)
        output, (h_n, c_n) = self.bilstm(emb)
        encoding1 = h_n[-2, :, :]
        encoding2 = h_n[-1, :, :]
        encoding = torch.cat([encoding1, encoding2], dim=-1)
        out = self.fc(encoding)
        out = F.log_softmax(out, dim=-1)
        return out
